import cv2

MIN_CONTOUR_AREA = 120 # 小于此面积不算作轮廓
FONT = cv2.FONT_HERSHEY_SIMPLEX
THRESH, MAX_VALUE = 230, 255
RED = (0, 0, 255)
BLACK = (0, 0, 0)
FONT_SCALE = 0.5
LINEWIDTH = 2


def get_type(cnt):
    epsilon = 0.01 * cv2.arcLength(cnt, True)
    approx = cv2.approxPolyDP(cnt, epsilon, True)
    edge_num = len(approx)
    if edge_num == 3:
        return "Tri"
    if edge_num == 4:
        return "Rect"
    if edge_num >= 5:
        return "Circ"
    return "Unknown"

# 传入彩色图， 输出处理好的图片
def process(img):
    # 预处理
    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    cv2.threshold(img_gray, THRESH, MAX_VALUE, cv2.THRESH_BINARY_INV, img_gray)


    contours, hierarchy = cv2.findContours(\
        img_gray, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
    

    # 筛选出符合条件的轮廓
    needed_contours = [contour for contour in contours\
        if cv2.contourArea(contour) > MIN_CONTOUR_AREA]
    cv2.drawContours(img, needed_contours, -1, RED)


    for contour in needed_contours:
        # 画矩形框
        x, y = contour[:, 0, 0], contour[:, 0, 1]
        start_point = [min(x), min(y)]
        end_point = [max(x), max(y)]
        cv2.rectangle(img, start_point, end_point, BLACK, LINE_WIDTH)
        
        # 标类型
        figure_type = get_type(contour)
        cv2.putText(img, figure_type, start_point, FONT, FONT_SCALE, BLACK)

    return img


if __name__ == "__main__":
    scr = cv2.imread("opencv.png")
    new_scr = process(scr)
    cv2.imshow("test", new_scr)
    cv2.waitKey()